from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

from ssl.sslserver import read_ecc_private_key_pem

# 服务端生成私钥和公钥
server_private_key = read_ecc_private_key_pem('./ecc-key.pem')
server_public_key = server_private_key.public_key()

# 客户端生成私钥和公钥
client_private_key = read_ecc_private_key_pem('./client-ecc-key.pem')
client_public_key = client_private_key.public_key()

# 客户端和服务端交换公钥
# 这里模拟客户端接收到服务端的公钥
received_server_public_key = server_public_key

# 服务端接收到客户端的公钥
received_client_public_key = client_public_key

# 客户端使用自己的私钥和服务端的公钥生成共享密钥
client_shared_secret = client_private_key.exchange(ec.ECDH(), received_server_public_key)

# 服务端使用自己的私钥和客户端的公钥生成共享密钥
server_shared_secret = server_private_key.exchange(ec.ECDH(), received_client_public_key)

# 对共享密钥进行派生，生成最终的对称密钥（可选步骤）
def derive_key(shared_secret):
    derived_key = HKDF(
        algorithm=hashes.SHA256(),
        length=48,
        salt=None,
        info=b'tls12 master secret',
    ).derive(shared_secret)
    return derived_key

client_derived_key = derive_key(client_shared_secret)
server_derived_key = derive_key(server_shared_secret)

# 验证客户端和服务端生成的派生密钥是否相同
if client_derived_key == server_derived_key:
    print("客户端和服务端生成的派生密钥相同，共享密钥生成成功！")
    print(f"派生密钥（十六进制）: {client_derived_key.hex()}")
else:
    print("客户端和服务端生成的派生密钥不同，共享密钥生成失败！")
    